package main

import (
	"container/heap"
	"fmt"
	"sync"
	"sync/atomic"
	"time"
)

type Task struct {
	id       int64
	execute  func()
	expire   time.Time
	idx      int
	canceled bool
}

type Timer struct {
	tasks    *TaskHeap
	lookup   map[int64]*Task
	mu       sync.Mutex
	addChan  chan *Task
	stopChan chan struct{}
	wg       sync.WaitGroup
	idGen    int64
}

type TaskHeap []*Task

// 实现五个函数构建堆
func (h TaskHeap) Len() int           { return len(h) }
func (h TaskHeap) Less(i, j int) bool { return h[i].expire.Before(h[j].expire) }
func (h TaskHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i]; h[i].idx = i; h[j].idx = j }

func (h *TaskHeap) Push(t interface{}) {
	n := len(*h)
	task := t.(*Task)
	task.idx = n
	*h = append(*h, task)
}

func (h *TaskHeap) Pop() interface{} {
	old := *h
	n := len(old)
	task := old[n-1]
	task.idx = -1
	*h = old[0 : n-1]
	return task
}

func NewTimer() *Timer {
	t := &Timer{
		tasks:    &TaskHeap{},
		lookup:   make(map[int64]*Task),
		addChan:  make(chan *Task),
		stopChan: make(chan struct{}),
	}
	heap.Init(t.tasks)
	t.wg.Add(1)
	go t.scheduler()
	return t
}

func (t *Timer) addTask(delay time.Duration, job func()) int64 {
	task := &Task{
		id:      atomic.AddInt64(&t.idGen, 1),
		execute: job,
		expire:  time.Now().Add(delay),
	}

	t.mu.Lock()
	t.lookup[task.id] = task
	t.mu.Unlock()

	t.addChan <- task

	return task.id
}

func (t *Timer) remove(id int64) {
	t.mu.Lock()
	defer t.mu.Unlock()

	if task, exists := t.lookup[id]; exists {
		task.canceled = true
		delete(t.lookup, id)
	}
}

func (t *Timer) scheduler() {
	defer t.wg.Done()

	var cur *time.Timer
	var timerChan <-chan time.Time

	for {
		t.mu.Lock()
		if t.tasks.Len() > 0 {
			for t.tasks.Len() > 0 {
				task := (*t.tasks)[0]
				if task.canceled { //已删除
					heap.Pop(t.tasks)
					delete(t.lookup, task.id)
				} else {
					break
				}
			}

			if t.tasks.Len() > 0 {
				duration := time.Until((*t.tasks)[0].expire)
				if cur != nil {
					cur.Stop()
				}
				cur = time.NewTimer(duration)
				timerChan = cur.C
			}
		}
		t.mu.Unlock()

		select {
		case <-t.stopChan:
			if cur != nil {
				cur.Stop()
			}
			return

		case task := <-t.addChan:
			t.mu.Lock()
			heap.Push(t.tasks, task)
			t.mu.Unlock()

		case <-timerChan:
			t.mu.Lock()
			if t.tasks.Len() == 0 {
				t.mu.Unlock()
				continue
			}

			task := heap.Pop(t.tasks).(*Task)
			delete(t.lookup, task.id)

			if !task.canceled {
				go task.execute() // 异步执行
			}
			t.mu.Unlock()
		}
	}
}

func (t *Timer) Stop() {
	close(t.stopChan)
	t.wg.Wait()
}

func main() {
	timer := NewTimer()
	defer timer.Stop()

	timer.addTask(2*time.Second, func() {
		fmt.Println("执行任务1")
	})

	id2 := timer.addTask(3*time.Second, func() {
		fmt.Println("执行任务2")
	})

	// 删除任务2
	timer.remove(id2)

	// 添加新任务
	timer.addTask(4*time.Second, func() {
		fmt.Println("执行任务3")
	})

	// 阻塞等待
	time.Sleep(5 * time.Second)
}
